# 导入必要的库
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.datasets import mnist
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense
from tensorflow.keras.utils import to_categorical

# 1. 加载MNIST数据集
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# 2. 数据预处理
# 归一化处理，使像素值在[0, 1]之间
x_train, x_test = x_train / 255.0, x_test / 255.0

# 调整数据形状为 (28, 28, 1)，适应CNN输入
x_train = x_train.reshape(-1, 28, 28, 1)
x_test = x_test.reshape(-1, 28, 28, 1)

# 对标签进行独热编码
y_train = to_categorical(y_train, 10)
y_test = to_categorical(y_test, 10)

# 3. 构建卷积神经网络（CNN）模型
model = Sequential([
    Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    MaxPooling2D(pool_size=(2, 2)),
    Conv2D(64, (3, 3), activation='relu'),
    MaxPooling2D(pool_size=(2, 2)),
    Flatten(),
    Dense(128, activation='relu'),
    Dense(10, activation='softmax')  # 输出10个类别
])

# 4. 编译模型
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 5. 训练模型
model.fit(x_train, y_train, epochs=5, batch_size=64, validation_split=0.2)

# 6. 在测试集上评估模型性能
test_loss, test_accuracy = model.evaluate(x_test, y_test)
print(f'Test accuracy: {test_accuracy}')

# 7. 使用模型进行预测
predictions = model.predict(x_test)

# 8. 可视化前10个预测结果
for i in range(10):
    plt.imshow(x_test[i].reshape(28, 28), cmap='gray')
    plt.title(f"Prediction: {np.argmax(predictions[i])}, Actual: {np.argmax(y_test[i])}")
    plt.show()
